// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/udf.h>

namespace kumo::pollux::functions {
    template<typename TExecCtx>
    struct SplitToMultiMapFunction {
        using RestMapType =
        melon::F14FastMap<std::string_view, std::vector<std::string_view> >;
        POLLUX_DEFINE_FUNCTION_TYPES(TExecCtx);

        void call(
            out_type<Map<Varchar, Array<Varchar> > > &out,
            const arg_type<Varchar> &input,
            const arg_type<Varchar> &entryDelimiter,
            const arg_type<Varchar> &keyValueDelimiter) {
            POLLUX_USER_CHECK(!entryDelimiter.empty(), "entryDelimiter is empty");
            POLLUX_USER_CHECK(!keyValueDelimiter.empty(), "keyValueDelimiter is empty");
            POLLUX_USER_CHECK_NE(
                entryDelimiter,
                keyValueDelimiter,
                "entryDelimiter and keyValueDelimiter must not be the same");

            if (input.empty()) {
                return;
            }

            callImpl(
                out,
                toStringView(input),
                toStringView(entryDelimiter),
                toStringView(keyValueDelimiter));
        }

        // Overloaded call function to handle UnknownValue cases
        template<typename TEntryDelimiter, typename TKeyValueDelimiter>
        void call(
            out_type<Map<Varchar, Array<Varchar> > > &out,
            const arg_type<Varchar> &input,
            const TEntryDelimiter &entryDelimiter,
            const TKeyValueDelimiter &keyValueDelimiter) {
            static_assert(
                std::is_same_v<TEntryDelimiter, arg_type<UnknownValue> > ||
                std::is_same_v<TKeyValueDelimiter, arg_type<UnknownValue> >);
            return;
        }

    private:
        static std::string_view toStringView(const arg_type<Varchar> &input) {
            return std::string_view(input.data(), input.size());
        }

        void callImpl(
            out_type<Map<Varchar, Array<Varchar> > > &out,
            std::string_view input,
            std::string_view entryDelimiter,
            std::string_view keyValueDelimiter) const {
            size_t pos = 0;

            RestMapType restMap;

            while (pos < input.size()) {
                const auto nextEntryPos = input.find(entryDelimiter, pos);

                // Process the current entry
                processEntry(
                    restMap,
                    // If no more delimiters are found, take the rest of the string as the
                    // next entry
                    nextEntryPos == std::string::npos
                        ? std::string_view(input.data() + pos, input.size() - pos)
                        : std::string_view(input.data() + pos, nextEntryPos - pos),
                    keyValueDelimiter);

                if (nextEntryPos == std::string::npos) {
                    break;
                }
                // Move to the next entry
                pos = nextEntryPos + entryDelimiter.size();
            }

            // Populate the output map
            out.reserve(restMap.size());
            for (const auto &[key, values]: restMap) {
                auto [keyWriter, valueWriter] = out.add_item();
                keyWriter.copy_from(key);
                valueWriter.copy_from(values);
            }
        }

        void processEntry(
            RestMapType &restMap,
            std::string_view entry,
            std::string_view keyValueDelimiter) const {
            const auto delimiterPos = entry.find(keyValueDelimiter);

            POLLUX_USER_CHECK_NE(
                delimiterPos,
                std::string::npos,
                "Key-value delimiter must appear exactly once in each entry. Bad input: '{}'",
                entry);

            const auto key = std::string_view(entry.data(), delimiterPos);

            const auto value = std::string_view(
                entry.data() + delimiterPos + keyValueDelimiter.size(),
                entry.size() - delimiterPos - keyValueDelimiter.size());

            // Validate that the value does not contain the delimiter
            POLLUX_USER_CHECK_EQ(
                value.find(keyValueDelimiter),
                std::string::npos,
                "Key-value delimiter must appear exactly once in each entry. Bad input: '{}'",
                entry);

            auto [it, _] = restMap.try_emplace(key);
            it->second.push_back(value);
        }
    };
} // namespace kumo::pollux::functions
